import sys
import argparse

from ag.config import Config
from ag.io import stdout_stream, print_verbose, stdout, stderr
from .__version__ import __version__


def ask(config: Config, message: list) -> str:
    try:
        if not config.onetime:
            completion = config.client.chat.completions.create(
                # 模型列表：https://help.aliyun.com/zh/model-studio/getting-started/models
                model=config.model.name,
                messages=message,
                stream=True,
                stream_options={"include_usage": config.verbose}
            )
            output = ""
            for chunk in completion:
                # 如果 stream_options.include_usage 为 True，则最后一个 chunk 的 choices 字段为空列表，需要跳过（可以通过chunk.usage获取 Token 使用量）
                if chunk.choices:
                    output += chunk.choices[0].delta.content
                    stdout_stream(chunk.choices[0].delta.content)
                    if chunk.choices[0].finish_reason == "stop":
                        stdout("\n")
                elif config.verbose:
                    print_verbose(chunk)
            return output
        else:
            if config.json:
                completion = config.client.chat.completions.create(
                    model=config.model.name,
                    messages=message,
                    stream=False,
                    response_format={"type": "json_object"},
                )
            else:
                completion = config.client.chat.completions.create(
                    model=config.model.name,
                    messages=message,
                    stream=False,
                )
            output = completion.choices[0].message.content
            stdout(output)
            if config.verbose:
                print_verbose(completion)
            return output
    except Exception as err:
        stderr(f"错误信息：{err}")
        stderr("请参考文档：https://help.aliyun.com/zh/model-studio/developer-reference/error-code")
        return ""


def main():
    parser = argparse.ArgumentParser(description="处理输入参数")
    parser.add_argument('-m', '--model',        type=str, default="qwen3:8b", help='[qwen3:8b] 指定模型名称')
    parser.add_argument('-f', '--file',         type=str,                     help='[        ] 指定要读取的文件路径')
    parser.add_argument('-o', '--onetime',      action='store_true',          help='[false   ] 关闭流式输入功能')
    parser.add_argument('-i', '--interactive',  action='store_true',          help='[false   ] 开启交互模式')
    parser.add_argument('-v', '--verbose',      action='store_true',          help='[false   ] 显示token使用情况')
    parser.add_argument('-r', '--remember',     action='store_true',          help='[false   ] 交互模式下保存每轮对话')
    parser.add_argument('-j', '--json',         action='store_true',          help='[false   ] 开启结构化输出模式(该模式非流式输出)')
    parser.add_argument('-p', '--prefix',       nargs="?", type=str,          help='[        ] 提示词前缀信息')
    parser.add_argument('--version',            action='store_true',          help='[false   ] 显示版本')
    parser.add_argument('prompt',               nargs="?", type=str,          help='[        ] 使用双引号括起来的提示信息')

    args = parser.parse_args()
    if args.version:
        print(__version__)
        sys.exit(0)

    config = Config(parser.parse_args())

    if config.interactive:
        message = []
        while True:
            talk = input(">>> ")
            if talk == "exit":
                break
            talk = config.prompt + talk
            if not args.remember:
                # 不记忆则每轮都重置message
                message = []
            message.append({'role': 'user', 'content': talk})
            output = ask(config, message)
            message.append({'role': 'system', 'content': output})
    else:
        if len(config.prompt) <= 1:
            print(__version__)
            sys.exit(0)
        message = [
            {'role': 'system', 'content': 'You are a helpful assistant.'},
            {'role': 'user', 'content': config.prompt}
        ]
        ask(config, message)


if __name__ == '__main__':
    main()
